import tensorflow as tf

"""unet with resnet50 pretrained"""

IMAGE_SIZE = 256

input_shape = (IMAGE_SIZE, IMAGE_SIZE, 3)


def Unet_model():
    encoder = tf.keras.applications.ResNet50(include_top=False,
                                             input_shape=input_shape,
                                             )

    skip_connection_names = ["input_1", "conv1_relu", "conv2_block3_out", "conv3_block4_out"]
    encoder_output = encoder.get_layer("conv4_block6_out").output

    f = [64, 128, 256, 512]
    x = encoder_output
    for i in range(1, len(skip_connection_names) + 1, 1):
        x_skip = encoder.get_layer(skip_connection_names[-i]).output
        x = tf.keras.layers.UpSampling2D((2, 2))(x)
        # x = tf.keras.layers.Conv2DTranspose(f[-i], 2, 2, padding='same')(x)
        x = tf.keras.layers.Concatenate()([x, x_skip])

        x = tf.keras.layers.Conv2D(f[-i], (3, 3), padding="same")(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Activation("relu")(x)

        x = tf.keras.layers.Conv2D(f[-i], (3, 3), padding="same")(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Activation("relu")(x)

    x = tf.keras.layers.Conv2D(1, (1, 1), padding="same")(x)
    x = tf.keras.layers.Activation("sigmoid")(x)

    model = tf.keras.Model(encoder.inputs, x)
    return model


if __name__ == '__main__':
    # input_shape = (256, 256, 3)
    # resnet50 = tf.keras.applications.ResNet50(include_top=False,
    #                                           input_shape=input_shape,
    #                                           )
    # resnet50.summary()
    m = Unet_model()
    m.summary()
